#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under theglim Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys
import os
import cProfile

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False
    else:
        PlotStuff=True
        import matplotlib
        from pylab import *
        from plot_aux import *
        if sys.argv[0][-7:] == '-nox.py':
            matplotlib.use('Agg') 



import scipy.io

from numpy import *

from corisco import Picture
import corisco_aux
from quaternion import Quat, random_quaternion

import scipy.optimize

import Image

set_printoptions(precision=7)



import filterSQP

######################################################################
## Functions used in the optimization

## Conic constraint
def val_c(x):
    return np.linalg.norm(x)
def grad_c(x):
    return x/np.linalg.norm(x)
def hess_c(x):
    nx = np.linalg.norm(x)
    return (nx**2 * np.identity(x.shape[0]) - np.outer(x,x)) / nx**3

## Target function
def val_f(x,*fargs):
    return corisco_aux.angle_error(x, fargs[0],fargs[1],fargs[2])
def grad_f(x,*fargs):
    return corisco_aux.angle_error_gradient(x, fargs[0],fargs[1],fargs[2])
def hess_f(x,*fargs):
    return corisco_aux.angle_error_hessian(x, fargs[0],fargs[1],fargs[2])
##
######################################################################


def utime():
    '''Pick user time from os to measure the code performance.'''
    return os.times()[0]

if __name__ == '__main__':
    tt_total_init = utime()
    if PlotStuff:
        rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load image and initialize pic object
    tt_init = utime()

    ## Sets filename from input argument
    if len(sys.argv) < 2:
        print sys.argv[0], '<filename> [grid_spacing] [gradient_threshold] [focal_distance] [initial_trials] [noise_scale] [smoothing_factor]'
        exit
        raise Exception('Insufficient number of parameters')

    ## Open image
    filename = sys.argv[1]
    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    ## Set all the job parameters (camera model, error function model,
    ## image processing parameters, etc)
    job_params = {}

    job_params['projection_model'] = 'pinhole'
    job_params['grid_spacing'] = 4
    job_params['gradient_threshold'] = 15.0
    job_params['focal_distance'] = imr.shape[1] * 32./36.
    job_params['principal_point'] = [imr.shape[1]*.5, imr.shape[0]*.5]
    job_params['noise_scale'] = 0.15
    job_params['smoothing_factor'] = 0.0

    job_params['do_robust'] = 1
    job_params['do_optimization'] = 1
    job_params['edge_detection_method'] = 2 ## Shigeru gradient filter

    job_params['initial_trials'] = 10000 ## Initial RANSAC trials
    job_params['optimization_tolerance'] = 1e-14

    

    if len(sys.argv) > 2:
        job_params['grid_spacing'] = int(sys.argv[2])
    if len(sys.argv) > 3:
        job_params['gradient_threshold'] = float(sys.argv[3])
    if len(sys.argv) > 4:
        job_params['focal_distance'] = float(sys.argv[4])
    if len(sys.argv) > 5:
        job_params['initial_trials'] = int(sys.argv[5])
    if len(sys.argv) > 6:
        job_params['noise_scale'] = float(sys.argv[6])
    if len(sys.argv) > 7:
        job_params['gaussian_smoothing_factor'] = float(sys.argv[7])

    job_params['fp_robust'] = [3.0, 1.0, job_params['noise_scale']]
    job_params['fp_optimization'] = [3.0, 1.0, job_params['noise_scale']]


    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], \
                          double(job_params["gaussian_smoothing_factor"]))

    ## Creates picture object
    if not 'projection_model' in job_params:
        raise "Missing camra model in job file"
    else:
        model = job_params['projection_model']

    if not (model in ['pinhole', 'harris', 'polar_equidistant', 'cylindrical_equidistant']):
        raise NotImplementedError

    elif model == 'pinhole':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        i_param = array([0.0, focal_distance, p_point[0], p_point[1]])

    elif model == 'harris':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        distortion = job_params['distortion_coefficient']
        i_param = array([2.0, focal_distance, p_point[0], p_point[1], distortion])
        
    elif model == 'polar_equidistant':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        i_param = array([3.0, focal_distance, p_point[0], p_point[1]])

    elif model == 'cylindrical_equidistant':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        i_param = array([4.0, focal_distance, p_point[0], p_point[1]])

    sqp_funcs = (val_c, grad_c, hess_c, val_f, grad_f, hess_f)
    pic = Picture(imr, i_param)


    tt_image_initialization = utime()-tt_init
    tt_directional_filter = 0.0
    ##
    ##################################################################

    ## Edgel extractor parameters
    gmethod = job_params['edge_detection_method']
    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold']

    do_robust = job_params['do_robust']
    initial_trials = job_params['initial_trials']

    do_optimization = job_params['do_optimization']
    optimization_tolerance = job_params['optimization_tolerance']


    ##################################################################
    ## Extract the edgels from the image using the grid mask
    tt_init = utime()
    pic.extract_edgels(gspc, glim, method=gmethod)
    tt_edgel_extraction = utime()-tt_init
    ##
    ##################################################################

    ##################################################################
    ## Caculate the edgel (interpretation plane) normals.
    pic.calculate_edgel_normals()
    ##
    ##################################################################

    ##################################################################
    ## Calculate Jacobians at each edgel
    tt_init = utime()
    if do_robust == 1:
        pic.calculate_edgel_jacobians()
    tt_jacobians = utime() - tt_init
    ##
    ##################################################################

    ##################################################################
    ## Calculate initial estimate. (First optimization step, based on
    ## RANSAC.)
    tt_init = utime()
    if do_robust == 1:
        fp = array(job_params["fp_robust"])
        qini,bpk_a,bpk_b,bpk_c = pic.random_search(initial_trials, fp)
        qini = qini.canonical()
    else:
        initial_trials = 0
    tt_initial_estimate = utime() - tt_init
    ##
    ##################################################################

    #################################################################
    ## Perform optimization procedure to find orientation even more
    ## accurately. (Second optimization step, based on FitlerSQP.)
    tt_init = utime()

    fp = array(job_params["fp_optimization"])
    args_f = (pic.edgels, i_param, fp)

    filterSQPout = filterSQP.filterSQP(qini.q,.0,1e-3,sqp_funcs,args_f,
                                       delta_tol=optimization_tolerance)

    xo, err, sqp_iters,Llam,Lrho = filterSQPout

    qopt = Quat(xo)

    tt_filtersqp = utime() - tt_init
    ##
    ##################################################################

    tt_total = utime() - tt_total_init

    if PlotStuff:
        figure(1)
        title('Extracted edgels')
        plot_edgels(gca(),pic)

        figure(11)
        title(u'Direções preditas com a orientação estimada')
        imshow(.1+pic.frame*.9/255.)
        pic.plot_vdirs(gca(), 31, qopt)
        axis([0,pic.Iwidth, pic.Iheight, 0])

    qini = qini.canonical()
    qopt = qopt.canonical()

    print filename,
    print 4*' % 10.8f'%(qini.q[0], qini.q[1], qini.q[2], qini.q[3]),
    print 4*' % 10.8f'%(qopt.q[0], qopt.q[1], qopt.q[2], qopt.q[3]),
    print pic.edgels.shape[0], initial_trials, sqp_iters,
    print 5*' %.2f'%(tt_edgel_extraction, tt_directional_filter,
                     tt_initial_estimate, tt_filtersqp, tt_total)




    #################################################################
    ## Generate the rectified output
    R = qopt.inverse().canonical().rot()

    os = 1.0
    outfd = 400.
    pp = array([
            [-outfd, -outfd, 2*outfd],
            [-outfd,  outfd, 2*outfd],
            [ outfd,  outfd, 2*outfd],
            [ outfd, -outfd, 2*outfd],
            ])

    qq = dot(pp,R)
    
    q = qq[:,:2]*pic.fd/qq[:,[2,2]]
    q += array([pic.middlex, pic.middley])

    aa = im.transform( ( int(2.0*outfd), int(2.0*outfd)),
                       Image.QUAD, q.ravel(), Image.BILINEAR )

    aa.save(filename + '-RECTIFIED.png')

